/*
 * @lc app=leetcode.cn id=337 lang=swift
 *
 * [337] 打家劫舍 III
 */

// @lc code=start
/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     public var val: Int
 *     public var left: TreeNode?
 *     public var right: TreeNode?
 *     public init() { self.val = 0; self.left = nil; self.right = nil; }
 *     public init(_ val: Int) { self.val = val; self.left = nil; self.right = nil; }
 *     public init(_ val: Int, _ left: TreeNode?, _ right: TreeNode?) {
 *         self.val = val
 *         self.left = left
 *         self.right = right
 *     }
 * }
 */
class Solution {
    func rob(_ root: TreeNode?) -> Int {
        // 返回一个2的数组
        // arr[0]: 不抢root的最大钱
        // arr[1]: 抢root的最大钱数
        func dp(_ root: TreeNode?) -> [Int] {
            if root == nil { return [0, 0] }

            let left = dp(root?.left)
            let right = dp(root?.right)

            // 抢，不能抢下家了
            let do_it = root!.val + left[0] + right[0]
            // 不抢，下家可抢可不抢
            let not_do = max(left[0], left[1]) + max(right[0], right[1])

            return [not_do, do_it]
        }

        let robs = dp(root)
        return max(robs[0], robs[1])
    }
}
// @lc code=end

